from flask import Flask, request, send_from_directory, jsonify
import json
from flask_socketio import SocketIO, send, emit
import os
from RAG import RAG


app = Flask(__name__)
socketio = SocketIO(app)
rag = RAG(embedding_model="BAAI/bge-large-zh-v1.5", llm_model="Alibaba/Qwen1.5-7B")


@app.route('/')
def home():
    return send_from_directory('dist', 'index.html')

@app.route('/<path:path>')
def static_file(path):
    return send_from_directory('dist', path)


@socketio.on('chat')
def chat(para):
    para = json.loads(para)
    order_id = para['orderId']
    history = para['history'] 
    query = para['query']
    file_names = para['fileNames']
    if len(file_names) > 0:  # 用户上传了文件，则召回文本
        targets = rag.recall_records(query=query, file_names=file_names)
        targets_content = "\n".join([x.text for x in targets])
        for item in rag.chat(query=query, history=history, targets_content=targets_content):
            item['orderId'] = order_id
            send(item)
        send({
            "orderId": order_id,
            "finish_reason": "recall",
            "value": [x.to_dict() for x in targets]
        })  
    else: # 用户没有上传了文件，仅实现基础LLM
        for item in rag.chat_no_files(query=query, history=history):
            item['orderId'] = order_id
            send(item)
         

@app.route('/upload', methods=['POST'])
def upload_files():
    files = request.files.getlist('files')  # 获取文件列表
    if not files:
        return jsonify({'error': 'No files found'})

    filenames = []
    for file in files:
        if file:
            filename = file.filename
            file.save(os.path.join('uploads', filename))
            rag.add_document(filename)
            filenames.append(filename)

    return jsonify({'message': f'{len(files)} files uploaded successfully', 'statistics': rag.db_discribe(), "upload": filenames})


@app.route('/feedback', methods=['POST'])
def feedback():
   # 获取JSON数据
    data = request.get_json()

    # 从数据中提取数组和字符串
    history = data.get('history')
    recalls = data.get('recalls')
    result = data.get('result')

    rag.add_feedback(history=history, recalls=recalls, result=result)  
    return jsonify({'message': 'Feedback received successfully'}) 


def ensure_directory_exists(directory_path):
    """
    检查指定的路径是否存在一个文件夹，如果不存在，则创建它。

    :param directory_path: str, 需要检查或创建的文件夹的路径。
    """
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)
        print(f"Directory created: {directory_path}")
    else:
        print(f"Directory already exists: {directory_path}")

        
if __name__ == '__main__':
    ensure_directory_exists('uploads')
    app.run(debug=False, host='0.0.0.0', port=5002)
